package com.fs.chatx.module.knowledge.service;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fs.chatx.module.chat.service.EmbeddingService;
import com.fs.chatx.module.knowledge.vo.DataItem;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.github.pigmesh.ai.deepseek.core.EmbeddingClient;
import io.github.pigmesh.ai.deepseek.core.embedding.Embedding;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.Float16Vec;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Service
public class MilvusOpService
{
    @Autowired
    private  MilvusClientV2 milvusClient;
    @Autowired
    private EmbeddingService embeddingService;

    @Value("${milvus.collectionName}")
    private String collectionName;
    @Autowired
    private Gson gson;


    @PostConstruct
    public void init()
    {
        try {
            createCollectionIfNotExists(collectionName,getModelDimension());
        } catch (Exception e) {
            e.printStackTrace();
           log.error("初始化milvus集合异常:{}",e.getMessage());
        }
    }

    public void createCollectionIfNotExists(String collectionName, int vectorDimension)
    {
        boolean exists = milvusClient.hasCollection(HasCollectionReq.builder()
                .collectionName(collectionName)
                .build());
        log.info("集合{},是否存在:{}",collectionName,exists);
        if (!exists) {

            CreateCollectionReq.CollectionSchema collectionSchema = CreateCollectionReq.CollectionSchema.builder()
                    .build();
            collectionSchema.addField(AddFieldReq.builder()
                    .fieldName("id")
                    .dataType(DataType.Int64)
                    .isPrimaryKey(true)
                    .autoID(true)
                    .build());

//            collectionSchema.addField(AddFieldReq.builder()
//                    .fieldName("original")
//                    .dataType(DataType.VarChar)
//                    .description("原始数据")
//                    .build());



            collectionSchema.addField(AddFieldReq.builder()
                    .fieldName("source")
                    .dataType(DataType.VarChar)
                    .description("来源")
                    .maxLength(500)
                    .build());

            collectionSchema.addField(AddFieldReq.builder()
                    .fieldName("question")
                    .dataType(DataType.VarChar)
                    .description("问题")
                    .maxLength(500)
                    .build());
            collectionSchema.addField(AddFieldReq.builder()
                    .fieldName("answer")
                    .dataType(DataType.VarChar)
                    .description("答案")
                    .build());


            collectionSchema.addField(AddFieldReq.builder()
                    .fieldName("embedding")
                    .dataType(DataType.FloatVector)
                    .dimension(vectorDimension)
                    .description("向量")
                    .build());
            //索引
            List<IndexParam> indexes = new ArrayList<>();
            indexes.add(IndexParam.builder()
                    .fieldName("embedding")
                    .indexType(IndexParam.IndexType.IVF_FLAT)
                    .metricType(IndexParam.MetricType.COSINE)
                    .build());

            CreateCollectionReq requestCreate = CreateCollectionReq.builder()
                    .collectionName(collectionName)
                    .collectionSchema(collectionSchema)
                    .indexParams(indexes)
                    .consistencyLevel(ConsistencyLevel.BOUNDED)
                    .build();
            milvusClient.createCollection(requestCreate);
            log.info("集合{}:Collection created sucess",collectionName);
        }
    }

    private int getModelDimension() {
        int  modelDimension=  embeddingService.embed("向量长度计算").size(); // 或动态值
        log.info("modelDimension:{}",modelDimension);
        return modelDimension;
    }

    public List<List<SearchResp.SearchResult>> searchList(String collectionName, String question, String vectorFieldName)
    {
        long begin = System.currentTimeMillis();
        List<Float> floatList = embeddingService.embed(question);
        List<String> outFields= Arrays.asList("question","answer");
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(Collections.singletonList(new FloatVec(floatList)))
//                .searchParams()
//                .annsField(vectorFieldName)
                .outputFields(outFields)
                .topK(5)
                .build();
        SearchResp searchResp = milvusClient.search(searchReq);
        long end = System.currentTimeMillis();
        long cost = (end - begin);
        System.out.println("Search time cost: " + cost + "ms");
        return  searchResp.getSearchResults();
    }


    public void importDataFromJson() throws IOException
    {
        List<DataItem> jsonDataSet = readJsonFile();
        int rowCount = jsonDataSet.size();
        int batchSize = 1000; // 每批插入1000条
        // 1、 提取所有需要嵌入的文本
        List<String> textsToEmbed = jsonDataSet.stream()
                .map(DataItem::getEmbeddingText) // 假设DataItem有getText()方法
                .collect(Collectors.toList());
        //2、 分批次生成嵌入向量
        long begin = System.currentTimeMillis();
        int embedBatchSize = 10;
        List<List<Float>> embeddings = new ArrayList<>();
        for (int i = 0; i < textsToEmbed.size(); i += embedBatchSize) {
            int end = Math.min(i + embedBatchSize, textsToEmbed.size());
            List<String> batchTexts = textsToEmbed.subList(i, end);
            List<Embedding> batchEmbeddings = embeddingService.batchEmbed(batchTexts);
            batchEmbeddings.forEach(e -> {
                embeddings.add(e.embedding());
            });
        }

        long endTime = System.currentTimeMillis();
        long cost = (endTime - begin);
        System.out.println("分批次生成嵌入向量 time cost: " + cost + "ms");

        // 3、 构建插入数据
        begin = System.currentTimeMillis();
        List<JsonObject> rows = new ArrayList<>(rowCount);
        Gson gson = new Gson();
        for (int j = 0; j < rowCount; j++) {
            JsonObject row = new JsonObject();
            row.add("embedding", gson.toJsonTree(embeddings.get(j)));
            row.addProperty("source",jsonDataSet.get(j).getSource());
            row.addProperty("question",jsonDataSet.get(j).getQuestion());
            row.addProperty("answer",jsonDataSet.get(j).getAnswer());

            rows.add(row);
        }
        // 4、分批次插入Milvus
        for (int i = 0; i < rowCount; i += batchSize) {
            int end = Math.min(i + batchSize, rowCount);
            List<JsonObject> batch = rows.subList(i, end);
            InsertResp resp = milvusClient.insert(InsertReq.builder()
                    .collectionName(collectionName)
                    .data(batch)
                    .build());
        }
         endTime = System.currentTimeMillis();
         cost = (endTime - begin);
        System.out.println("分批次插入Milvus向量 time cost: " + cost + "ms");


    }


    public List<DataItem> readJsonFile() throws IOException
    {
        List<DataItem> lines = new ArrayList<>();
        try
        {
            Resource jsonResource = new ClassPathResource("data/v00034.json");
            BufferedReader reader = new BufferedReader(
                    new InputStreamReader(jsonResource.getInputStream(), StandardCharsets.UTF_8)) ;
            String line;
            while ((line = reader.readLine()) != null) {
                JSONObject  lineObject = JSON.parseObject(line);
                // 解析QA字段
                String qaText =lineObject.get("qa").toString();
                int questionIndex = qaText.indexOf("提问：");
                int answerIndex = qaText.indexOf("回答:");
                String question = qaText.substring(questionIndex + 3, answerIndex).trim();
                String answer = qaText.substring(answerIndex + 3).trim();
                // 解析来源
                String source = lineObject.getJSONObject("source").getOrDefault("参考","").toString();
                lines.add(DataItem.builder().embeddingText(question)
                                .source(source)
                                .question(question)
                                .answer(answer)
                        .build());
               // log.info("question:{}  answer:{}",question,answer);
            }

        }catch (IOException e) {
            e.printStackTrace();
            throw e;
        }

        return lines;
    }


}
